'''
handle all the trained image to abstract graph and transform
'''
import os
import json
from util.graph import Graph, transform_graph, GraphTransformer
import argparse
from util.vessel_seg import build_vessels

def get_configs():
    '''
    参数：
    data_path: 数据集路径
    transform_number: 进行增强的数量
    vascular_height: 血管图像高度 default450
    vascular_width: 血管图像宽度 default=60
    '''
    parser = argparse.ArgumentParser(description="Configuration for ROP dataset processing and model training")
    parser.add_argument('--data_path', type=str, default='../Dataset/infantImages', help="Path to the dataset")
    parser.add_argument('--transform_number', type=int, default=10, help="Number of transformations")
    parser.add_argument('--vascular_height', type=int, default=450, help="Height of the vascular image")
    parser.add_argument('--vascular_width', type=int, default=600, help="Width of the vascular image")
    
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = get_configs()
    data_path = args.data_path
    transform_number = args.transform_number
    vascular_height= args.vascular_height
    vascular_width= args.vascular_width
    
    
    # 读取数据字典
    with open(os.path.join(data_path, 'annotations.json'), 'r') as f:
        data_dict = json.load(f)
    # 创建图转换器
    graph_transformer = GraphTransformer(image_size=(vascular_width, vascular_height))
    
        
    split_name = 'all'
    with open(os.path.join(data_path, 'split', f'{split_name}.json'), 'r') as f:
        used_image = json.load(f)
    graph_record = {}
    for image_name in used_image['train']:
        vascular_path=data_dict[image_name]['vascular_path']
        x,y=data_dict[image_name]['optic_disc_pred']['position']
        vessels= build_vessels(vascular_path, (x, y))
        
        # 创建图结构
        graph = Graph(vascular_height,vascular_width)
        
        graph.build_from_vessel(vessels)
        
        record={'transform_number':transform_number,
                'data':[]}
        # 保存原始图
        record['data'].append(graph._2json())
        for i in range(transform_number):
            # 图转换
            new_graph,missing_num,node_map = transform_graph(graph, graph_transformer)
            new_graph_json = new_graph._2json()
            new_graph_json['missing_num'] = missing_num
            new_graph_json['node_map'] = node_map
            record['data'].append(new_graph_json)
        
        graph_record[image_name]=record
    with open(os.path.join(data_path, 'graph_record.json'), 'w') as f:
        json.dump(graph_record, f)
        
    print('Graphs saved successfully.')
            